extern "C" __device__ 
inline float vec3_dot(float* v1, float* v2)
{
    return v1[0] * v2[0] + v1[1] * v2[1] + v1[2] * v2[2];
}
extern "C" __device__ 
inline float vec4_dot(float* v1, float* v2)
{
    return v1[0] * v2[0] + v1[1] * v2[1] + v1[2] * v2[2] + v1[3] * v2[3];
}

extern "C" __device__ 
void mat4_mul_vec4(float* mat4, float* vec4, float* out) {
    float x0 = vec4_dot(mat4, vec4);
    float x1 = vec4_dot(mat4 + 4, vec4);
    float x2 = vec4_dot(mat4 + 8, vec4);
    float x3 = vec4_dot(mat4 + 12, vec4);
    out[0] = x0;
    out[1] = x1;
    out[2] = x2;
    out[3] = x3;
}
extern "C" __global__ 
void mat4_transform(float* vertices, float* transformed, float* mat, unsigned int vert_num)
{
    int vert_idx = blockIdx.x * blockDim.x + threadIdx.x; 
    if (vert_idx < vert_num) 
    {
        mat4_mul_vec4(mat, vertices + vert_idx * 4, transformed + vert_idx * 4);
    }
}
extern "C" __global__ 
void mat4_transform_inplace(float* vertices, float* mat, unsigned int vert_num)
{
    mat4_transform(vertices, vertices, mat, vert_num);
}
extern "C" __global__ 
void vertex_ndc_transform(float* vertices, unsigned int vert_num)
{
    int vert_idx = blockIdx.x * blockDim.x + threadIdx.x; 
    if (vert_idx < vert_num) 
    {
        float w = vertices[vert_idx * 4 + 3];
        vertices[vert_idx * 4] = vertices[vert_idx * 4] / w;
        vertices[vert_idx * 4 + 1] = vertices[vert_idx * 4 + 1] / w;
        vertices[vert_idx * 4 + 2] = vertices[vert_idx * 4 + 2] / w;
        vertices[vert_idx * 4 + 3] = 1;
    }
}

extern "C" __device__ 
inline float signed_triangle_area(int ax, int ay, int bx, int by, int cx, int cy) 
{
    return 0.5 * ((by-ay)*(bx+ax) + (cy-by)*(cx+bx) + (ay-cy)*(ax+cx));
}
extern "C" __device__ 
inline int min2(int a, int b) 
{
    return max(min(a, b), 0);
}
extern "C" __device__ 
inline int max2(int a, int b, int bound) 
{
    return min(max(a, b), bound);
}
extern "C" __device__ 
inline int min3(int a, int b, int c) 
{
    return max(min(min(a, b), c), 0);
}
extern "C" __device__ 
inline int max3(int a, int b, int c, int bound) 
{
    return min(max(max(a, b), c), bound);
}
extern "C" __device__ 
inline void cross(float* a, float* b, float* ret) {
    ret[0] = a[1] * b[2] - a[2] * b[1];
    ret[1] = a[2] * b[0] - a[0] * b[2];
    ret[2] = a[0] * b[1] - a[1] * b[0]; 
}
extern "C" __device__ 
inline void norm(float* x) {
    float len = sqrt(vec3_dot(x, x));
    x[0] = x[0] / len;
    x[1] = x[1] / len;
    x[2] = x[2] / len;
}
extern "C" __device__ 
inline void reflect(float* n, float* l, float* ret) {
    float cos2 = vec3_dot(n, l) * 2;
    ret[0] = n[0] * cos2 - l[0];
    ret[1] = n[1] * cos2 - l[1];
    ret[2] = n[2] * cos2 - l[2];
}

struct RasterizationKernelParam
{
    float* model_view_mat;
    float* normal_mat;

    float* model_verts;
    unsigned int* model_face_vrts;

    float* model_norms;
    unsigned int* model_face_norm;
    
    float* model_tex;
    unsigned int* model_face_tex;

    unsigned long face_num;

    unsigned char* normal_map;
    unsigned long normal_map_width;
    unsigned long normal_map_height;

    unsigned char* diffuse_map;
    unsigned long diffuse_map_width;
    unsigned long diffuse_map_height;

    char* frame_buffer;
    float* z_buffer;
    unsigned long fb_width;
    unsigned long fb_height;
};
extern "C" __global__ 
void rasterization(char* param)
{
    struct RasterizationKernelParam* params = (struct RasterizationKernelParam*)param;

    int face_idx = blockIdx.x * blockDim.x + threadIdx.x; 
    if (face_idx < params->face_num)
    {
        int v1_idx = face_idx * 3;
        int v2_idx = face_idx * 3 + 1;
        int v3_idx = face_idx * 3 + 2;

        float* v1 = params->model_verts + params->model_face_vrts[v1_idx] * 4;
        float* v2 = params->model_verts + params->model_face_vrts[v2_idx] * 4;
        float* v3 = params->model_verts + params->model_face_vrts[v3_idx] * 4;

        float* n1 = params->model_norms + params->model_face_norm[v1_idx] * 4;
        float* n2 = params->model_norms + params->model_face_norm[v2_idx] * 4;
        float* n3 = params->model_norms + params->model_face_norm[v3_idx] * 4;

        float* t1 = params->model_tex + params->model_face_tex[v1_idx] * 2;
        float* t2 = params->model_tex + params->model_face_tex[v2_idx] * 2;
        float* t3 = params->model_tex + params->model_face_tex[v3_idx] * 2;

        int width = params->fb_width;
        int height = params->fb_height;

        bool ok = v1[0] >= 0 && v1[0] < width && v1[1] >= 0 && v1[1] < height;
        ok = ok || (v2[0] >= 0 && v2[0] < width && v2[1] >= 0 && v2[1] < height);
        ok = ok || (v3[0] > -1 && v3[0] < 1 && v3[1] > -1 && v3[1] < 1);

        if (ok)
        {
            int ax = v1[0], ay = v1[1];
            float az = v1[2];
            int bx = v2[0], by = v2[1];
            float bz = v2[2];
            int cx = v3[0], cy = v3[1];
            float cz = v3[2];

            int bbminx = min3(ax, bx, cx); // bounding box for the triangle
            int bbminy = min3(ay, by, cy); // defined by its top left and bottom right corners
            int bbmaxx = max3(ax, bx, cx, width - 1);
            int bbmaxy = max3(ay, by, cy, height - 1);
            float total_area = signed_triangle_area(ax, ay, bx, by, cx, cy);
            if (total_area < 1) return;

            float l[4];
            l[0] = 1;
            l[1] = 1;
            l[2] = 1;
            l[3] = 0;
            mat4_mul_vec4(params->model_view_mat, l, l);
            norm(l);

            for (int x = bbminx; x <= bbmaxx; x++) {
                for (int y = bbminy; y <= bbmaxy; y++) {
                    float alpha = signed_triangle_area(x, y, bx, by, cx, cy) / total_area;
                    float beta  = signed_triangle_area(x, y, cx, cy, ax, ay) / total_area;
                    float gamma = signed_triangle_area(x, y, ax, ay, bx, by) / total_area;

                    if (alpha >= 0 && beta >= 0 && gamma >= 0) {
                        float z = alpha * az + beta * bz + gamma * cz;
                        int fb_idx = x + y * width;

                        if (z < params->z_buffer[fb_idx])
                        {
                            params->z_buffer[fb_idx] = z;

                            uchar4 color;
                            // pixel shader
                            {
                                float ambient = 0.4;

                                float tx = alpha * t1[0] + beta * t2[0] + gamma * t3[0];
                                float ty = alpha * t1[1] + beta * t2[1] + gamma * t3[1];

                                int normal_idx_x = min2(int(tx * params->normal_map_width), params->normal_map_width - 1);
                                int normal_idx_y = min2(int(ty * params->normal_map_height), params->normal_map_height - 1);
                                unsigned char* normal = params->normal_map + (normal_idx_x + normal_idx_y * params->normal_map_width) * 4;

                                float n[4];
                                n[0] = float(normal[0]) * 2. / 255. - 1.;
                                n[1] = float(normal[1]) * 2. / 255. - 1.;
                                n[2] = float(normal[2]) * 2. / 255. - 1.;
                                n[3] = 0.;
                                mat4_mul_vec4(params->normal_mat, n, n);
                                norm(n);
                                float diffuse = vec3_dot(l, n);
                                if (diffuse < 0) diffuse = 0;

                                float r[3];
                                reflect(n, l, r);
                                float r_dot_v = r[2] > 0 ? r[2] : 0;
                                float specular = pow(r_dot_v, 35);

                                float coe = ambient + 1.0 * diffuse + 1.5 * specular;

                                int diffuse_idx_x = min2(int(tx * params->diffuse_map_width), params->diffuse_map_width - 1);
                                int diffuse_idx_y = min2(int((1 - ty) * params->diffuse_map_height), params->diffuse_map_height - 1);
                                unsigned char* diffuse_color = params->diffuse_map + (diffuse_idx_x + diffuse_idx_y * params->diffuse_map_width) * 4;
                            
                                color.x = min2(coe * diffuse_color[0], 255);
                                color.y = min2(coe * diffuse_color[1], 255);
                                color.z = min2(coe * diffuse_color[2], 255);
                                color.w = 255;
                            }
                            
                            uchar4 *framebuffer = (uchar4 *)(params->frame_buffer);
                            framebuffer[fb_idx] = color;
                        }    
                    }
                }
            }
        }
    }
}